package query

import (
	"bufio"
	"fmt"
	"strconv"
	"strings"
)

// IluvatarQuery Iluvatar GPU查询实现
type IluvatarQuery struct{}

func init() {
	RegisterQueryer("iluvatar", func() GPUQuery { return &IluvatarQuery{} })
}

func (q *IluvatarQuery) Query() ([]GPUInfo, error) {
	output, err := runCommand("/usr/local/corex/bin/ixsmi",
		"--query-gpu=index,utilization.gpu,memory.used,memory.total",
		"--format=csv,noheader,nounits")
	if err != nil {
		return nil, fmt.Errorf("执行ixsmi失败: %v", err)
	}

	return parseIluvatarOutput(output)
}

// parseIluvatarOutput 解析Iluvatar GPU输出
func parseIluvatarOutput(output string) ([]GPUInfo, error) {
	var gpus []GPUInfo
	scanner := bufio.NewScanner(strings.NewReader(output))

	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Split(line, ",")
		if len(fields) < 4 {
			continue
		}

		// 解析每个字段
		index, err1 := strconv.Atoi(strings.TrimSpace(fields[0]))
		util, err2 := strconv.Atoi(strings.TrimSpace(fields[1]))
		used, err3 := strconv.Atoi(strings.TrimSpace(fields[2]))
		total, err4 := strconv.Atoi(strings.TrimSpace(fields[3]))

		if err1 != nil || err2 != nil || err3 != nil || err4 != nil {
			continue
		}

		gpus = append(gpus, GPUInfo{
			Index:       index,
			Utilization: util,
			MemoryUsed:  used,
			MemoryTotal: total,
		})
	}

	if len(gpus) == 0 {
		return nil, fmt.Errorf("未找到有效的GPU信息")
	}

	return gpus, nil
}
